import xml.etree.ElementTree as ET
from colorama import Fore, Style

def scenario_analyzer(scenario_url: str, start_frame: int, end_frame: int):
    root = ET.ElementTree(file=scenario_url).getroot()
    num_frames = end_frame - start_frame + 1
    num_frame_vehicles = 0
    num_tasks = 0
    full_path_tasks = 0
    num_unique_vehicles = set()
    num_full_path_vehicles = set()
    num_full_path_frames = 0

    sum_speed = 0
    for step_i in range(len(root)):
        num_frame_vehicles += len(root[step_i])
        for vehicle in root[step_i]:
            num_tasks += len(vehicle)
            sum_speed += float(vehicle.attrib["speed"])
            if (vehicle.attrib["id"] not in num_unique_vehicles) and (float(vehicle.attrib["x"]) < 15 or float(vehicle.attrib["x"]) > 1185):
                num_full_path_vehicles.add(vehicle.attrib["id"])
            if vehicle.attrib["id"] in num_full_path_vehicles:
                full_path_tasks += len(vehicle)
                num_full_path_frames += 1
            num_unique_vehicles.add(vehicle.attrib["id"])
    print(Fore.GREEN + f"共有 {num_frames:d} 帧" + Style.RESET_ALL)
    print(Fore.GREEN + f"平均每帧中有 {num_frame_vehicles / num_frames:.1f} 辆车" + Style.RESET_ALL)
    print(Fore.GREEN + f"共有 {len(num_unique_vehicles):d} 辆车" + Style.RESET_ALL)
    print(Fore.GREEN + f"平均每辆车速度为 {sum_speed/num_frame_vehicles:.1f} m/s" + Style.RESET_ALL)

    print(Fore.YELLOW + f"完整行程车辆 平均每辆车在场景中的停留时间为 {(num_full_path_frames / len(num_full_path_vehicles)) / 10:.1f} 秒" + Style.RESET_ALL)
    print(Fore.YELLOW + f"完整行程车辆 平均每辆车产生 {full_path_tasks / len(num_full_path_vehicles):.1f} 个任务" + Style.RESET_ALL)


if __name__ == '__main__':
    url = "../scenario_xml/scenario_highway_vehsPerHour=500_lam=0.35.xml"
    # url = "../scenario_xml/scenario_gridNetwork_p=1.0_lam=0.25.xml"
    scenario_analyzer(scenario_url=url,
                      start_frame=0,
                      end_frame=999)
